import akshare as ak
import pandas as pd
import datetime as dt
import logging,time,asyncio
from django.utils import timezone
from datetime import datetime, timedelta
from asgiref.sync import sync_to_async
from .models import StockMinuteData, IndexMinuteData, DataUpdateRecord,TradeDate,StockBasic
from django.db import transaction
logger = logging.getLogger(__name__)
# 创建异步的数据库操作函数
update_data_record = sync_to_async(DataUpdateRecord.update_record)


def get_start_date(is_index=False):
    """
    获取分时数据更新的起始日期
    规则：取(5个交易日前日期, 历史数据最大日期+1天)的最大值
    """
    # 获取当天（含）之前的5个交易日
    trade_dates = TradeDate.objects.filter(
            is_open=1,
            cal_date__lte=datetime.now().date().strftime('%Y%m%d')
        ).order_by('-cal_date').values_list('cal_date', flat=True)[:5]
    trade_dates = list(trade_dates)
    print(f'trade_dates  xxxx: {trade_dates}')#格式20250630
    early_date = datetime.strptime(trade_dates[-1],'%Y%m%d').date()#取最后一个，而不是第一个
    if is_index:
        record = DataUpdateRecord.objects.filter(data_type='index_minute').order_by('-last_update_date').first()
    else:
        record = DataUpdateRecord.objects.filter(data_type='stock_minute').order_by('-last_update_date').first()
    if record:#如果有记录，就取记录的第二天与前5个交易日中日期最小值
        start_date = max(record.last_update_date + timedelta(days=1),early_date)
    else:#如果没有记录，就取5个交易日前日期的最小值
        start_date = early_date
    #返回date格式
    return start_date



def count_trading_days(start_date, end_date):
    """
    计算两个日期之间的交易日数量
    :param start_date: 开始日期，datetime.date对象或START_KDATA_DATE常量
    :param end_date: 结束日期，datetime.date对象
    :return: 交易日数量
    """
    if not start_date or not end_date:
        return 0
    
    # 如果start_date是START_KDATA_DATE常量，直接使用它（因为它已经是datetime.date对象）
    
    if start_date > end_date:#起始日期和结束日期相同时可能为一天
        return 0
    
    try:
        # 将日期转换为字符串格式YYYYMMDD
        start_str = start_date.strftime('%Y%m%d')
        end_str = end_date.strftime('%Y%m%d')
        
        # 查询交易日历
        trading_days = TradeDate.objects.filter(
            cal_date__gte=start_str,
            cal_date__lte=end_str,
            is_open=True
        ).count()
        
        return trading_days
    except Exception as e:
        return 0

get_start_date_async = sync_to_async(get_start_date)
count_trading_days_async = sync_to_async(count_trading_days)

class AkshareAPI:
    def __init__(self):
        self.progress = 0
        self.status = ""

    async def update_all_data(self,start_date=None, end_date=None, websocket=None):
        """
        异步更新所有akshare数据
        :param data_types: 需要更新的数据类型列表，如果为None则更新所有数据
        """

        data_types = ['index_minute','stock_minute']
        
        results = {}
        print(f'update_all_data  ssssss')
        for data_type in data_types:
            self.status = f"正在更新 {data_type} 数据..."
            self.progress = 0
            try:
                if data_type == 'stock_minute':
                    # continue#暂时先不更新各个个股的分时数据
                    results[data_type] = await self.update_stock_minute_data(websocket=websocket)
                elif data_type == 'index_minute':
                    results[data_type] = await self.update_index_minute_data()
                    
            except Exception as e:
                print(f"更新 {data_type} 数据失败: {e}")
                results[data_type] = False
            
            self.progress = 100
        
        self.status = "数据更新完成"
        return results

   
    

    #还要支持单个更新数据
    async def update_stock_minute_data(self, stock_codes=None,websocket=None):
        """
        异步更新个股分时数据
        :param stock_codes: 需要更新的股票代码列表，如果为None则更新所有股票
        """
        # 获取智能计算的起始日期
        end_date = timezone.now().date()
        start_date = await get_start_date_async(is_index=False)
        print(f'start_date: {start_date} end_date: {end_date}')
        if start_date > end_date:
            print(f'没有需要更新的数据,start_date: {start_date} end_date: {end_date}')
            return False
        
        # 如果没有指定股票代码，可以从StockBasic模型中获取所有股票代码
        if stock_codes is None:
            stock_codes = await sync_to_async(list)(StockBasic.objects.exclude(name__contains='ST').values_list('ts_code', flat=True))
        
        total = len(stock_codes)
        updated = False
        current_step = 0
        total_steps = total
        batch_size = 100
        for i in range(0, total, batch_size):
            batch_stock_codes = stock_codes[i:i + batch_size]
            batch_df = pd.DataFrame()
            for j, code in enumerate(batch_stock_codes):
                self.progress = int((i*batch_size + j + 1) / total * 100)
                if websocket:
                    current_step += 1
                    await websocket.send_progress(
                        current=current_step,
                        total=total_steps,
                        message=f"正在更新股票分时数据 stock_code = {code}"
                )
                df = await self.get_stock_minute_data(code, start_date.strftime('%Y-%m-%d'))
                if df is not None and not df.empty:
                    batch_df = pd.concat([batch_df, df])
                else:
                    print(f"获取个股分时数据失败: {code}")
                    return
            if not batch_df.empty:
                if await self.save_stock_minute_data(batch_df):
                    updated = True
        
        if updated:
            await sync_to_async(DataUpdateRecord.update_record)('stock_minute', end_date)
        return updated

    async def update_index_minute_data(self, index_codes=None):
        """
        异步更新指数分时数据
        :param index_codes: 需要更新的指数代码列表，如果为None则更新主要指数
        """
        # 默认更新主要指数
        if index_codes is None:
            index_codes = ['sh000001', 'sz399001', 'sz399006', 'sz399300']
        
        # 获取智能计算的起始日期
        end_date = timezone.now().date()
        start_date = await get_start_date_async(is_index=True)#
        print(f'start_date: {start_date} end_date: {end_date}')
        if start_date > end_date:
            print(f'没有需要更新的数据,start_date: {start_date} end_date: {end_date}')
            return False
        
        total = len(index_codes)
        updated = False
        
        for i, code in enumerate(index_codes):
            self.progress = int((i + 1) / total * 100)
            df = await self.get_index_minute_data(code, start_date, end_date)
            if df is not None and not df.empty:
                if await self.save_index_minute_data(df):
                    updated = True
        
        if updated:
            await sync_to_async(DataUpdateRecord.update_record)('index_minute', end_date)
        return updated


    async def get_stock_minute_data(self,stock_code, start_date=None):
        """
        获取个股分时数据
        :param stock_code: 股票代码(带市场前缀，如'000001.SZ')
        :param start_date: 开始日期(可选)
        :param end_date: 结束日期(可选)
        :return: DataFrame
        """
        if not start_date:
            start_date = (timezone.now() - timedelta(days=5)).strftime('%Y%m%d')
        
        try:
            # 使用stock_zh_a_hist_min_em获取历史分时数据
            df = ak.stock_zh_a_hist_min_em(
                symbol=stock_code[:6],
                period='1',
                start_date=start_date,#只需要开始日期，这样最多默认取5天数据
                adjust=''
            )
            
            # 转换数据格式以匹配现有结构
            df['ts_code'] = stock_code
            df = df.rename(columns={
                '时间': 'trade_time',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '成交额': 'amount',
                '均价': 'avg_price',
            })
            
            return df[['ts_code', 'trade_time', 'open', 'high', 'low', 'close', 'volume', 'amount', 'avg_price']]
        except Exception as e:
            print(f"获取个股分时数据失败: {e}")
            return None
    #现在start_date和end_date都是date格式，而非字符串
    async def get_index_minute_data(self,index_code, start_date=None, end_date=None):
        """
        获取指数分时数据
        :param index_code: 指数代码(带市场前缀，如'sh000001')
        :param start_date: 开始日期(可选)
        :param end_date: 结束日期(可选)
        :return: DataFrame
        """
        
        if not start_date:
            start_date = timezone.now().date() - timedelta(days=5)
        if not end_date:
            end_date = timezone.now().date()
        trade_date_num = await count_trading_days_async(start_date, end_date)
        print(f'trade_date_num: {trade_date_num}')
        
        try:
            # 使用stock_zh_a_minute获取分时数据,每天只有239个数据，9:30和14:59的数据缺失
            df = ak.stock_zh_a_minute(symbol=index_code, period='1')#分时数据目前免费的只有新浪和东财，但只有新浪接口支持指数
            df = df[-239*trade_date_num:]
            # 确保trade_time是datetime类型
            df['trade_time'] = pd.to_datetime(df['day'])
            df['ts_code'] = index_code
            # start_date_str = start_date.strftime('%Y%m%d')
            # end_date_str = (end_date + timedelta(days=1)).strftime('%Y%m%d')#要多算一天
            # start_dt = datetime.strptime(start_date_str, '%Y%m%d')  # 转换为datetime类型
            # end_dt = datetime.strptime(end_date_str, '%Y%m%d')  # 转换为datetime类型
            # 过滤指定日期范围的数据
            # df = df[(df['trade_time'] >= start_dt) & (df['trade_time'] <= end_dt)]
            #插入一列amount
            df['amount'] = 0#指数的成交量，不怎么用。只为了保持接口一致（与个股）
            return df[['ts_code', 'trade_time', 'open', 'high', 'low', 'close', 'volume','amount']]
        except Exception as e:
            print(f"获取指数分时数据失败: {e}")
            return None

    @sync_to_async
    def save_stock_minute_data(self,df):
        """
        保存分时数据到数据库
        :param df: 包含分时数据的DataFrame
        
        """
        if df is None or df.empty:
            return False
        
        model = StockMinuteData
        records = []
        with transaction.atomic():
            for _, row in df.iterrows():
                records.append(model(
                    ts_code=row['ts_code'],
                    trade_time=row['trade_time'],
                    open=row['open'],
                    high=row['high'],
                    low=row['low'],
                    close=row['close'],
                    vol=row['volume'],
                    amount=row['amount'],
                    avg_price=row['avg_price']
                ))
        
            try:
                model.objects.bulk_create(records, ignore_conflicts=True)
                return True
            except Exception as e:
                print(f"保存分时数据失败: {e}")
                return False    
    @sync_to_async
    def save_index_minute_data(self,df):
        """
        保存分时数据到数据库
        :param df: 包含分时数据的DataFrame
        """
        if df is None or df.empty:
            return False
        
        model = IndexMinuteData
        records = []
        with transaction.atomic():
            for _, row in df.iterrows():
                records.append(model(
                    ts_code=row['ts_code'],
                    trade_time=row['trade_time'],
                    open=row['open'],
                    high=row['high'],
                    low=row['low'],
                    close=row['close'],
                    vol=row['volume'],
                    amount=row['amount']
                ))
            
            try:
                model.objects.bulk_create(records, ignore_conflicts=True)
                return True
            except Exception as e:
                print(f"保存指数分时数据失败: {e}")
                return False